"""Add retry fields to scheduled_task and custom_task

Revision ID: 987654321abc
Revises: 123456789abc
Create Date: 2025-07-21 22:30:00.000000

"""
from alembic import op
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = '987654321abc'
down_revision = '123456789abc'
branch_labels = None
depends_on = None


def upgrade():
    # 为scheduled_task表添加重试相关字段
    with op.batch_alter_table('scheduled_task', schema=None) as batch_op:
        batch_op.add_column(sa.Column('max_retries', sa.Integer(), nullable=True, server_default='3'))
        batch_op.add_column(sa.Column('retry_delay', sa.Integer(), nullable=True, server_default='300'))
        batch_op.add_column(sa.Column('retry_count', sa.Integer(), nullable=True, server_default='0'))
        batch_op.add_column(sa.Column('last_error', sa.Text(), nullable=True))
    
    # 为custom_task表添加重试相关字段
    with op.batch_alter_table('custom_task', schema=None) as batch_op:
        batch_op.add_column(sa.Column('max_retries', sa.Integer(), nullable=True, server_default='3'))
        batch_op.add_column(sa.Column('retry_delay', sa.Integer(), nullable=True, server_default='300'))
        batch_op.add_column(sa.Column('retry_count', sa.Integer(), nullable=True, server_default='0'))
        batch_op.add_column(sa.Column('last_error', sa.Text(), nullable=True))


def downgrade():
    # 从scheduled_task表删除重试相关字段
    with op.batch_alter_table('scheduled_task', schema=None) as batch_op:
        batch_op.drop_column('last_error')
        batch_op.drop_column('retry_count')
        batch_op.drop_column('retry_delay')
        batch_op.drop_column('max_retries')
    
    # 从custom_task表删除重试相关字段
    with op.batch_alter_table('custom_task', schema=None) as batch_op:
        batch_op.drop_column('last_error')
        batch_op.drop_column('retry_count')
        batch_op.drop_column('retry_delay')
        batch_op.drop_column('max_retries')